package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"ai-rag/config"
	"ai-rag/workflow"

	"github.com/gin-contrib/cors"
	"github.com/gin-gonic/gin"
	"github.com/joho/godotenv"
	"github.com/patrickmn/go-cache"
	"github.com/sirupsen/logrus"
)

// 缓存用于存储查询结果
var queryCache *cache.Cache

func main() {
	// 解析命令行参数
	configPath := flag.String("config", "config.toml", "配置文件路径")
	mode := flag.String("mode", "web", "运行模式: web, index, query")
	filePath := flag.String("file", "", "要索引的文件路径")
	query := flag.String("q", "", "要查询的问题")
	flag.Parse()

	// 加载环境变量
	if err := godotenv.Load(); err != nil {
		logrus.Warn("未找到.env文件，将使用默认配置")
	}

	// 加载配置
	cfg, err := config.LoadConfig(*configPath)
	if err != nil {
		logrus.Fatalf("加载配置失败: %v", err)
	}

	// 设置日志级别
	setupLogger(cfg.Service.LogLevel)

	// 注意：Redis功能暂时禁用，因为依赖问题
	if cfg.Service.Redis.Enable {
		logrus.Warn("Redis配置已启用，但由于依赖问题暂时禁用")
	}

	// 初始化缓存
	if cfg.Service.Cache.Enable {
		queryCache = cache.New(time.Duration(cfg.Service.Cache.Expiration)*time.Second, 10*time.Minute)
		logrus.Info("查询缓存已启用")
	}

	// 初始化RAG系统
	ragSystem, err := workflow.NewRAGSystem(cfg)
	if err != nil {
		logrus.Fatalf("初始化RAG系统失败: %v", err)
	}
	defer func() {
		if err := ragSystem.Close(); err != nil {
			logrus.Errorf("关闭RAG系统失败: %v", err)
		}
	}()

	// 根据运行模式执行不同的操作
	switch *mode {
	case "web":
		// 启动Web服务
		startWebServer(cfg, ragSystem)
	case "index":
		// 索引单个文件
		if *filePath == "" {
			logrus.Fatal("请提供要索引的文件路径")
		}
		_, err = ragSystem.IndexDocument(context.Background(), *filePath)
		if err != nil {
			logrus.Fatalf("索引文件失败: %v", err)
		}
		logrus.Info("文件索引成功")
	case "query":
		// 执行单个查询
		if *query == "" {
			logrus.Fatal("请提供查询内容")
		}
		response, err := ragSystem.AnswerQuery(context.Background(), *query)
		if err != nil {
			logrus.Fatalf("查询失败: %v", err)
		}
		fmt.Printf("回答: %s\n", response)
	default:
		logrus.Fatalf("未知的运行模式: %s", *mode)
	}
}

// setupLogger 设置日志级别
func setupLogger(logLevel string) {
	level, err := logrus.ParseLevel(logLevel)
	if err != nil {
		logrus.Warnf("无效的日志级别: %s, 使用默认级别info", logLevel)
		level = logrus.InfoLevel
	}
	logrus.SetLevel(level)
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp: true,
	})
}

// startWebServer 启动Web服务
// initRedisClient 初始化Redis客户端
// 注意：此函数暂时禁用，因为依赖问题
func initRedisClient(cfg *config.Config) {
	// Redis功能暂时禁用
	logrus.Warn("Redis初始化函数暂时禁用，因为依赖问题")
}

func startWebServer(cfg *config.Config, ragSystem *workflow.RAGSystem) {
	// 设置Gin模式
	if cfg.Service.LogLevel == "debug" {
		gin.SetMode(gin.DebugMode)
	} else {
		gin.SetMode(gin.ReleaseMode)
	}

	// 创建Gin引擎
	r := gin.New()
	r.Use(gin.Logger())
	r.Use(gin.Recovery())

	// 配置CORS
	if cfg.Service.EnableCORS {
		r.Use(cors.New(cors.Config{
			AllowOrigins:     []string{"*"},
			AllowMethods:     []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
			AllowHeaders:     []string{"Origin", "Content-Type", "Authorization"},
			ExposeHeaders:    []string{"Content-Length"},
			AllowCredentials: true,
			MaxAge:           12 * time.Hour,
		}))
	}

	// API路由
	api := r.Group("/api")
	{
		// 索引文档
		api.POST("/index", handleIndexDocument(ragSystem))

		// 查询文档
		api.POST("/query", handleQuery(ragSystem, cfg))

		// 健康检查
		api.GET("/health", func(c *gin.Context) {
			c.JSON(http.StatusOK, gin.H{
				"status":       "healthy",
				"redis_status": "disabled_due_to_dependency",
				"config": gin.H{
					"embedding_type": cfg.Embedding.Type,
					"indexer_type":   cfg.Indexer.Type,
				},
			})
		})
	}

	// 静态文件服务
	r.Static("/static", "./static")

	// 首页
	r.GET("/", func(c *gin.Context) {
		c.JSON(http.StatusOK, gin.H{
			"message": "RAG系统API服务运行中",
			"version": "1.0.0",
		})
	})

	// 创建HTTP服务器
	srv := &http.Server{
		Addr:    fmt.Sprintf(":%d", cfg.Service.Port),
		Handler: r,
	}

	// 优雅关闭
	go func() {
		// 等待中断信号
		quit := make(chan os.Signal, 1)
		signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
		<-quit
		logrus.Info("正在关闭服务器...")

		// 设置5秒的超时时间
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()

		if err := srv.Shutdown(ctx); err != nil {
			logrus.Fatal("服务器强制关闭:", err)
		}

		logrus.Info("服务器已优雅关闭")
	}()

	// 启动服务器
	logrus.Infof("RAG系统Web服务启动在 http://localhost:%d", cfg.Service.Port)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		logrus.Fatalf("启动服务器失败: %v", err)
	}
}

// handleIndexDocument 处理文档索引请求
func handleIndexDocument(ragSystem *workflow.RAGSystem) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req struct {
			FilePath string `json:"file_path" binding:"required"`
		}

		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数错误: " + err.Error()})
			return
		}

		// 验证文件是否存在
		if _, err := os.Stat(req.FilePath); os.IsNotExist(err) {
			c.JSON(http.StatusNotFound, gin.H{"error": "文件不存在"})
			return
		}

		// 索引文档
		startTime := time.Now()
		_, err := ragSystem.IndexDocument(c.Request.Context(), req.FilePath)
		if err != nil {
			logrus.WithError(err).Error("索引文档失败")
			c.JSON(http.StatusInternalServerError, gin.H{"error": "索引文档失败: " + err.Error()})
			return
		}

		duration := time.Since(startTime)
		logrus.WithField("file_path", req.FilePath).WithField("duration", duration).Info("文档索引成功")

		c.JSON(http.StatusOK, gin.H{
			"message":   "文档索引成功",
			"file_path": req.FilePath,
			"duration":  duration.Milliseconds(),
		})
	}
}

// handleQuery 处理查询请求
func handleQuery(ragSystem *workflow.RAGSystem, cfg *config.Config) gin.HandlerFunc {
	return func(c *gin.Context) {
		var req struct {
			Query   string                 `json:"query" binding:"required"`
			Context map[string]interface{} `json:"context,omitempty"`
		}

		if err := c.ShouldBindJSON(&req); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"error": "请求参数错误: " + err.Error()})
			return
		}

		// 检查缓存
		if cfg.Service.Cache.Enable && queryCache != nil {
			// 使用本地内存缓存
			if cached, found := queryCache.Get(req.Query); found {
				logrus.WithField("query", req.Query).Info("从内存缓存返回查询结果")
				c.JSON(http.StatusOK, gin.H{
					"query":      req.Query,
					"answer":     cached.(string),
					"from_cache": true,
					"cache_type": "memory",
				})
				return
			}
		}

		// 记录配置信息
		logrus.WithFields(logrus.Fields{
			"embedding_server":  cfg.Embedding.ServerURL,
			"indexer_type":      cfg.Indexer.Type,
			"ark_api_enabled":   cfg.Service.ArkAPI.Enabled,
			"agent_api_enabled": cfg.Service.AgentAPI.Enabled,
		}).Debug("查询配置信息")

		// 执行查询
		startTime := time.Now()
		response, err := ragSystem.AnswerQuery(c.Request.Context(), req.Query)
		if err != nil {
			logrus.WithError(err).WithField("query", req.Query).Error("查询失败")
			c.JSON(http.StatusInternalServerError, gin.H{"error": "查询失败: " + err.Error()})
			return
		}

		duration := time.Since(startTime)
		logrus.WithField("query", req.Query).WithField("duration", duration).Info("查询成功")

		// 缓存结果
		if cfg.Service.Cache.Enable && queryCache != nil {
			// 使用本地内存缓存
			queryCache.Set(req.Query, response, cache.DefaultExpiration)
		}

		c.JSON(http.StatusOK, gin.H{
			"query":    req.Query,
			"answer":   response,
			"duration": duration.Milliseconds(),
		})
	}
}
